#include <arpa/inet.h>
#include <errno.h>
#include <memory.h>
#include <netinet/in.h>
#include <openssl/aes.h>
#include <openssl/bn.h>
#include <openssl/dh.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#include <iostream>

#include <google/protobuf/message.h>
#include "login.pb.h"
using namespace google::protobuf;
using namespace std;

#define MAXLINE 4096

// g++ -I ~/Downloads/openssl-1.0.2l/include -I /usr/local/include -L/usr/local/lib -lcrypto -lssl -lprotobuf -o client main.cpp login.pb.cc

//填充需要加密数据长度为16的倍数
char* pKCS5Padding(unsigned char* src, int len, int blockSize, int* outLen) {
  int padding = blockSize - (len % blockSize);
  char* str = new char[len + padding];
  memset(str, (char)padding, len + padding);
  memcpy(str, src, len);
  *outLen = len + padding;
  return str;
}

//去除加密时需要填充的数据
char* pKCS5UnPadding(unsigned char* src, int len, int* outLen) {
  int unpadding = int(src[len - 1]);

  if (unpadding > len) {
    cout << "error decrypt data, unpadding: " << unpadding << endl;
    return NULL;
  }

  char* str = new char[len - unpadding];
  memcpy(str, src, len - unpadding);
  *outLen = len - unpadding;
  return str;
}

char* AesEncrypt(unsigned char* key, char* src, int len, int* outLen) {
  AES_KEY aeskey;
  AES_set_encrypt_key(key, 256, &aeskey);
  //
  char* padding = pKCS5Padding((unsigned char*)src, len, 16, outLen);
  //
  char* out = new char[*outLen];
  unsigned char iv[AES_BLOCK_SIZE * 4] = {0};
  AES_cbc_encrypt((unsigned char*)padding, (unsigned char*)out, *outLen,
                  &aeskey, iv, AES_ENCRYPT);

  delete[] padding;
  return out;
}

char* AesDecrypt(unsigned char* key, char* src, int len, int* outLen) {
  AES_KEY aeskey;
  AES_set_decrypt_key(key, 256, &aeskey);
  //
  char* de = new char[len];

  unsigned char iv[AES_BLOCK_SIZE * 4] = {0};
  AES_cbc_encrypt((unsigned char*)src, (unsigned char*)de, len, &aeskey, iv,
                  AES_DECRYPT);

  char* buff = pKCS5UnPadding((unsigned char*)de, len, outLen);
  delete[] de;
  return buff;
}

void heartbeat(int fd, unsigned char* key, unsigned short cmd,
               unsigned short seq) {
  char buf1[4];
  cmd = htons(cmd);
  seq = htons(seq);
  memcpy(buf1, &cmd, 2);
  memcpy(buf1 + 2, &seq, 2);

  int enLen = 0;
  char* enstr = AesEncrypt(key, buf1, 4, &enLen);

  char buf3[2];
  unsigned short len1 = htons(enLen);
  memcpy(buf3, &len1, 2);

  send(fd, buf3, 2, 0);
  int len = send(fd, enstr, enLen, 0);
  delete[] enstr;

  char buf2[18];
  len = recv(fd, buf2, 18, 0);

  memcpy(&len, buf2, 2);
  int deLen = 0;
  char* destr = AesDecrypt(key, buf2 + 2, len, &deLen);

  int cmd1, seq1, ret1;
  memcpy(&cmd1, destr, 2);
  memcpy(&seq1, destr + 2, 2);
  memcpy(&ret1, destr + 4, 2);

  cmd1 = ntohs(cmd1);
  seq1 = ntohs(seq1);
  ret1 = ntohs(ret1);

  cout << "heartbeat result，cmd：" << cmd1 << ", seq: " << seq1
       << ", ret: " << ret1 << endl;
}

void register_player(int fd, unsigned char* key, unsigned short cmd,
                     unsigned short seq) {
  login::LoginReq req;
  login::LoginRsp rsp;
  req.set_username("guangbo");
  req.set_password("123123");
  int reqLen = req.ByteSize();

  char* buf1 = new char[4 + reqLen];
  cmd = htons(cmd);
  seq = htons(seq);
  memcpy(buf1, &cmd, 2);
  memcpy(buf1 + 2, &seq, 2);
  memcpy(buf1 + 4, req.SerializeAsString().data(), reqLen);

  int enLen = 0;
  char* enstr = AesEncrypt(key, buf1, 4 + reqLen, &enLen);

  char buf3[2];
  unsigned short len1 = htons(enLen);
  memcpy(buf3, &len1, 2);

  send(fd, buf3, 2, 0);
  int len = send(fd, enstr, enLen, 0);
  delete[] enstr;
  delete[] buf1;

  char buf2[1024];
  len = recv(fd, buf2, 1024, 0);

  memcpy(&len, buf2, 2);
  len = ntohs(len);
  int deLen = 0;
  char* destr = AesDecrypt(key, buf2 + 2, len, &deLen);

  int cmd1, seq1, ret1;
  memcpy(&cmd1, destr, 2);
  memcpy(&seq1, destr + 2, 2);
  memcpy(&ret1, destr + 4, 2);

  cmd1 = ntohs(cmd1);
  seq1 = ntohs(seq1);
  ret1 = ntohs(ret1);

  rsp.ParseFromArray(destr + 6, deLen - 4);

  cout << "regsiter result，cmd：" << cmd1 << ", seq: " << seq1
       << ", ret: " << ret1 << ", token: " << rsp.token() << endl;

  delete[] destr;
}

int connectServer() {
  int sockfd, n;
  char recvline[4096], sendline[4096];

  struct sockaddr_in servaddr;

  if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
    cout << "create socket error, " << strerror(errno) << endl;
    exit(0);
  }

  memset(&servaddr, 0, sizeof(servaddr));
  servaddr.sin_family = AF_INET;
  servaddr.sin_port = htons(10001);
  if (inet_pton(AF_INET, "127.0.0.1", &servaddr.sin_addr) <= 0) {
    cout << "inet_pton error for" << endl;
    exit(0);
  }

  if (connect(sockfd, (struct sockaddr*)&servaddr, sizeof(servaddr)) < 0) {
    cout << "connect error, " << strerror(errno) << endl;
    exit(0);
  }
  return sockfd;
}

int main() {
  char buf[MAXLINE];
  int rec_len;

  //先初始化dh数据
  DH* d1 = DH_new();
  int ret = DH_generate_parameters_ex(d1, 64 * 8, DH_GENERATOR_5, NULL);
  ret = DH_generate_key(d1);
  if (ret != 1) {
    cout << "DH_generate_key error" << endl;
    return 1;
  }

  int len = BN_bn2bin(d1->p, (unsigned char*)buf);
  len += BN_bn2bin(d1->g, (unsigned char*)buf + len);
  len += BN_bn2bin(d1->pub_key, (unsigned char*)buf + len);

  //连接服务器
  int sockfd = connectServer();

  //接受32字节随机字符串，用于最后密钥验证
  char random[34] = {0};
  if ((rec_len = recv(sockfd, random, 32, 0)) == -1) {
    cout << "recv error" << endl;
    exit(1);
  }

  //发送64字节p，1字节g，64字节客户端公钥
  len = send(sockfd, buf, len, 0);

  //接受64字节服务端公钥
  len = recv(sockfd, buf, 64, 0);

  BIGNUM* server_key = BN_new();
  BN_bin2bn((unsigned char*)buf, 64, server_key);

  //生成密钥，只取前面32字节即可
  unsigned char sharekey1[128] = {0};
  unsigned char key[32] = {0};
  len = DH_compute_key(sharekey1, server_key, d1);
  memcpy(key, sharekey1, 32);

  //加密签名获取的随机字符串，发送给服务端
  int outLen = 0;
  char* enrandom = AesEncrypt(key, random, 32, &outLen);
  len = send(sockfd, enrandom, outLen, 0);
  delete[] enrandom;

  cout << "handshake success" << endl;

  unsigned short seq = 1;
  while (true) {
    //一直发送注册
    register_player(sockfd, key, 10001, seq++);
    sleep(10);
  }
}